import torchaudio

from common.config import conf
from common.tts import TTS
import sys

from cosyvoice.cli.cosyvoice import CosyVoice2
from cosyvoice.utils.file_utils import load_wav

sys.path.insert(0, r'C:\tools\projects\python\role-bot\common\Matcha-TTS')


class CosyVoiceTTS(TTS):
    def __init__(self):
        super().__init__()
        config = conf()
        # SPEAKER
        speaker = config['speaker_audio']
        cosyvoice_model = config['cosyvoice_model']
        self.speaker_prompt = config['speaker_prompt']
        self.prompt_speaker_16k = load_wav(speaker, 16000)
        self.cosyvoice = CosyVoice2(cosyvoice_model)
        self.sample_rate = self.cosyvoice.sample_rate

    def text_to_voice(self, text):
        for _, j in enumerate(self.cosyvoice.inference_zero_shot(
                text, self.speaker_prompt, self.prompt_speaker_16k, stream=False)):
            return j['tts_speech']

    def text_to_voice_stream(self, text):
        for j in self.cosyvoice.inference_zero_shot(
                text, self.speaker_prompt, self.prompt_speaker_16k, stream=True):
            yield j['tts_speech']


if __name__ == '__main__':
    tts = CosyVoiceTTS()
    print(tts.sample_rate)
    torchaudio.save('zero_shot.wav',
                    tts.text_to_voice(
                        "收到好友从远方寄来的生日礼物，那份意外的惊喜与深深的祝福让我心中充满了甜蜜的快乐，笑容如花儿般绽放。"),
                    tts.sample_rate)
